import numpy as np
import torch
import torch.nn as nn
from models.nms_utils import grasps_distance_gpu 


class ProposalTargetLayer(nn.Module):
    def __init__(self, roi_sampler_cfg):
        super().__init__()
        self.roi_sampler_cfg = roi_sampler_cfg

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                rois: (B, num_rois, 10 + C)
                roi_scores: (B, num_rois)
                gt_grasps: (B, N, 10 + C + 1)
        Returns:
            batch_dict:
                rois: (B, M, 10 + C)
                gt_of_rois: (B, M, 10 + C)
                gt_iou_of_rois: (B, M)
                roi_scores: (B, M)
                reg_valid_mask: (B, M)
                rcnn_cls_labels: (B, M)
        """
        batch_rois, batch_gt_of_rois, batch_roi_ious, batch_roi_scores, sampled_inds = self.sample_rois_for_rcnn(
            batch_dict=batch_dict
        )
        # regression valid mask
        reg_valid_mask = (batch_roi_ious < self.roi_sampler_cfg.REG_FG_THRESH).long()

        # classification label
        if self.roi_sampler_cfg.CLS_SCORE_TYPE == 'cls':
            batch_cls_labels = (batch_roi_ious < self.roi_sampler_cfg.FG_THRESH).long()
            ignore_mask = (batch_roi_ious < self.roi_sampler_cfg.BG_THRESH) & \
                          (batch_roi_ious > self.roi_sampler_cfg.FG_THRESH)
            batch_cls_labels[ignore_mask > 0] = -1
        elif self.roi_sampler_cfg.CLS_SCORE_TYPE == 'roi_iou':
            raise NotImplementedError
            #iou_bg_thresh = self.roi_sampler_cfg.CLS_BG_THRESH
            #iou_fg_thresh = self.roi_sampler_cfg.CLS_FG_THRESH
            #fg_mask = batch_roi_ious > iou_fg_thresh
            #bg_mask = batch_roi_ious < iou_bg_thresh
            #interval_mask = (fg_mask == 0) & (bg_mask == 0)

            #batch_cls_labels = (fg_mask > 0).float()
            #batch_cls_labels[interval_mask] = \
            #    (batch_roi_ious[interval_mask] - iou_bg_thresh) / (iou_fg_thresh - iou_bg_thresh)
        else:
            raise NotImplementedError

        targets_dict = {'rois': batch_rois, 'gt_of_rois': batch_gt_of_rois, 
                        'gt_iou_of_rois': batch_roi_ious,
                        'roi_scores': batch_roi_scores,
                        'reg_valid_mask': reg_valid_mask,
                        'rcnn_cls_labels': batch_cls_labels, 
                        'sampled_inds': sampled_inds}
        # for debug_vis
        batch_dict['sampled_inds'] = sampled_inds

        return targets_dict

    def sample_rois_for_rcnn(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                rois: (B, num_rois, 10 + C)
                roi_scores: (B, num_rois)
                gt_grasps: (B, N, 10 + C)
        Returns:
        """
        batch_size = batch_dict['batch_size']
        rois = batch_dict['rois']
        roi_scores = batch_dict['roi_scores']
        #roi_labels = batch_dict['roi_labels']
        gt_grasps = batch_dict['gt_grasps']

        code_size = rois.shape[-1]
        batch_rois = rois.new_zeros(batch_size, self.roi_sampler_cfg.ROI_PER_IMAGE, code_size)
        batch_gt_of_rois = rois.new_zeros(batch_size, self.roi_sampler_cfg.ROI_PER_IMAGE, code_size+1)
        batch_roi_ious = rois.new_zeros(batch_size, self.roi_sampler_cfg.ROI_PER_IMAGE)
        batch_roi_scores = rois.new_zeros(batch_size, self.roi_sampler_cfg.ROI_PER_IMAGE)

        start = 0
        length = batch_dict['fg_list']
        for index in range(batch_size):
            cur_roi, cur_roi_scores = \
                rois[index], roi_scores[index]
            if batch_size > 1:
                end = start+length[index]
                cur_gt = gt_grasps[start:end]
                start += length[index]
            else:
                cur_gt = gt_grasps[start:start+length.item()]
            cur_gt = cur_gt.new_zeros((1, cur_gt.shape[1])) if len(cur_gt) == 0 else cur_gt

            iou3d = grasps_distance_gpu(cur_roi[:, 0:10].contiguous(), cur_gt[:, 0:10].contiguous())  # (M, N)

            min_distances, gt_assignment = torch.min(iou3d, dim=1)

            sampled_inds = self.subsample_rois(min_distances=min_distances)
            # for debug_vis
            batch_dict['target_sampled_inds'] = sampled_inds

            batch_rois[index] = cur_roi[sampled_inds]
            #batch_roi_labels[index] = cur_roi_labels[sampled_inds]
            batch_roi_ious[index] = min_distances[sampled_inds]
            batch_roi_scores[index] = cur_roi_scores[sampled_inds]
            batch_gt_of_rois[index] = cur_gt[gt_assignment[sampled_inds]]

        return batch_rois, batch_gt_of_rois, batch_roi_ious, batch_roi_scores, sampled_inds

    def subsample_rois(self, min_distances):
        # sample fg, easy_bg, hard_bg
        fg_rois_per_image = int(np.round(self.roi_sampler_cfg.FG_RATIO * self.roi_sampler_cfg.ROI_PER_IMAGE))
        #fg_thresh = min(self.roi_sampler_cfg.REG_FG_THRESH, self.roi_sampler_cfg.CLS_FG_THRESH)
        fg_thresh = self.roi_sampler_cfg.FG_THRESH

        fg_inds = ((min_distances <= fg_thresh)).nonzero().view(-1)
        easy_bg_inds = ((min_distances > self.roi_sampler_cfg.BG_THRESH_LO)).nonzero().view(-1)
        hard_bg_inds = ((min_distances > self.roi_sampler_cfg.FG_THRESH) &
                (min_distances <= self.roi_sampler_cfg.BG_THRESH_LO)).nonzero().view(-1)

        fg_num_rois = fg_inds.numel()
        bg_num_rois = hard_bg_inds.numel() + easy_bg_inds.numel()
        # set seed for debug
        #np.random.seed(0)
        if fg_num_rois > 0 and bg_num_rois > 0:
            # sampling fg
            fg_rois_per_this_image = min(fg_rois_per_image, fg_num_rois)

            rand_num = torch.from_numpy(np.random.permutation(fg_num_rois)).type_as(min_distances).long()
            fg_inds = fg_inds[rand_num[:fg_rois_per_this_image]]

            # sampling bg
            bg_rois_per_this_image = self.roi_sampler_cfg.ROI_PER_IMAGE - fg_rois_per_this_image
            bg_inds = self.sample_bg_inds(
                hard_bg_inds, easy_bg_inds, bg_rois_per_this_image, self.roi_sampler_cfg.HARD_BG_RATIO
            )

        elif fg_num_rois > 0 and bg_num_rois == 0:
            # sampling fg
            rand_num = np.floor(np.random.rand(self.roi_sampler_cfg.ROI_PER_IMAGE) * fg_num_rois)
            rand_num = torch.from_numpy(rand_num).type_as(min_distances).long()
            fg_inds = fg_inds[rand_num]
            bg_inds = []

        elif bg_num_rois > 0 and fg_num_rois == 0:
            # sampling bg
            bg_rois_per_this_image = self.roi_sampler_cfg.ROI_PER_IMAGE
            bg_inds = self.sample_bg_inds(
                hard_bg_inds, easy_bg_inds, bg_rois_per_this_image, self.roi_sampler_cfg.HARD_BG_RATIO
            )
        else:
            print('mindistances:(min=%f, max=%f)' % (min_distances.min().item(), min_distances.max().item()))
            print('ERROR: FG=%d, BG=%d' % (fg_num_rois, bg_num_rois))
            raise NotImplementedError
        
        if len(fg_inds) == 0:
            sampled_inds = bg_inds
        elif len(bg_inds) == 0:
            sampled_inds = fg_inds
        else:
            sampled_inds = torch.cat((fg_inds, bg_inds), dim=0)

        #print('forground num is : ', fg_inds.numel(), 'background num is : ', bg_inds.numel())
        return sampled_inds

    @staticmethod
    def sample_bg_inds(hard_bg_inds, easy_bg_inds, bg_rois_per_this_image, hard_bg_ratio):
        if hard_bg_inds.numel() > 0 and easy_bg_inds.numel() > 0:
            hard_bg_rois_num = min(int(bg_rois_per_this_image * hard_bg_ratio), len(hard_bg_inds))
            easy_bg_rois_num = bg_rois_per_this_image - hard_bg_rois_num

            # sampling hard bg
            rand_idx = torch.randint(low=0, high=hard_bg_inds.numel(), size=(hard_bg_rois_num,)).long()
            hard_bg_inds = hard_bg_inds[rand_idx]

            # sampling easy bg
            rand_idx = torch.randint(low=0, high=easy_bg_inds.numel(), size=(easy_bg_rois_num,)).long()
            easy_bg_inds = easy_bg_inds[rand_idx]

            bg_inds = torch.cat([hard_bg_inds, easy_bg_inds], dim=0)
        elif hard_bg_inds.numel() > 0 and easy_bg_inds.numel() == 0:
            hard_bg_rois_num = bg_rois_per_this_image
            # sampling hard bg
            rand_idx = torch.randint(low=0, high=hard_bg_inds.numel(), size=(hard_bg_rois_num,)).long()
            bg_inds = hard_bg_inds[rand_idx]
        elif hard_bg_inds.numel() == 0 and easy_bg_inds.numel() > 0:
            easy_bg_rois_num = bg_rois_per_this_image
            # sampling easy bg
            rand_idx = torch.randint(low=0, high=easy_bg_inds.numel(), size=(easy_bg_rois_num,)).long()
            bg_inds = easy_bg_inds[rand_idx]
        else:
            raise NotImplementedError

        return bg_inds

